# # 导入所需要的库
# import cv2
# import numpy as np
#
#
# # 定义保存图片函数
# # image:要保存的图片名字
# # addr；图片地址与相片名字的前部分
# # num: 相片，名字的后缀。int 类型
# def save_image(image, addr, num):
#     address = addr + str(num) + '.jpg'
#     cv2.imwrite(address, image)
#
#
# # 读取视频文件
# videoCapture = cv2.VideoCapture("2.mp4")
# # 通过摄像头的方式
# # videoCapture=cv2.VideoCapture(1)
#
# # 读帧
# success, frame = videoCapture.read()
# i = 0
# while success:
#     i = i + 1
#     save_image(frame, './output/image', m)
#     if success:
#         print('save image:', i)
#
# # 导入所需要的库
# import cv2
# import numpy as np
#
#
# # 定义保存图片函数
# # image:要保存的图片名字
# # addr；图片地址与相片名字的前部分
# # num: 相片，名字的后缀。int 类型
# def save_image(image, addr, num):
#     address = addr + str(num) + '.jpg'
#     cv2.imwrite(address, image)
#
#
# # 读取视频文件
# videoCapture = cv2.VideoCapture("2.mp4")
# # 通过摄像头的方式
# # videoCapture=cv2.VideoCapture(1)
#
# # 读帧
# success, frame = videoCapture.read()
# i = 0
# timeF = 12
# j = 0
# while success:
#     i = i + 1
#     if (i % timeF == 0):
#         j = j + 1
#         save_image(frame, './output/image', j)
#         print('save image:', i)
#     success, frame = videoCapture.read()

# coding=utf-8
import os
import random
import math
import time

import cv2
import uuid
import threading

import numpy as np
import torch
import os
from PIL import Image, ImageDraw
from facenet_pytorch.models.mtcnn import MTCNN, extract_face
from matplotlib import pyplot as plt

mtcnn = MTCNN(min_face_size=200, thresholds=[0.6, 0.7, 0.7], post_process=True, device='cuda')

face_cascade = cv2.CascadeClassifier(r'D:\anaconda\envs\ml\Lib\site-packages\cv2\data\haarcascade_frontalface_alt.xml')


def progress(percent, width=50):
    '''进度打印功能'''
    if percent >= 100:
        percent = 100

    show_str = ('[%%-%ds]' % width) % (int(width * percent / 100) * "#")  # 字符串拼接的嵌套使用
    print('\r%s %d%%' % (show_str, percent), end='')


def get_face(img):
    """这个算法的背景色与人脸差不多，性能很差"""
    s = 200  # 脸的最小的大小
    faces = face_cascade.detectMultiScale(img, 1.2, 9, minSize=(s, s), maxSize=(500, 500))

    if len(faces) == 1:
        for (x, y, w, h) in faces:
            padding = int(round(w * 0.06))
            # 扩大图片，可根据坐标调整
            X = int(x) - padding
            W = min(int(x + w + 2 * padding), img.shape[1])
            if X < 0:
                W -= X
                X = 0
            Y = int(y) - padding
            H = min(int(y + h + 2 * padding), img.shape[0])
            if Y < 0:
                H -= Y
                Y = 0
            return cv2.resize(img[Y:H, X:W], (W - X, H - Y))
    else:
        return None


def image_generator_count(path_, target=100):
    """target 目标取几张"""
    vc = cv2.VideoCapture(path_)
    # fps = vc.get(cv2.CAP_PROP_FPS)
    fcount = vc.get(cv2.CAP_PROP_FRAME_COUNT)
    count = 0
    trycount = 1
    # 随机大法
    while count < target:
        randSequence = random.sample(range(0, math.floor(fcount) - 1), round((target - count) * 1.1))
        randSequence = sorted(randSequence)
        for index in randSequence:
            vc.set(cv2.CAP_PROP_POS_FRAMES, index)
            rval, frame = vc.read()
            trycount += 1
            if rval:
                # 转换图片的格式
                frame = Image.fromarray(frame)
                r, g, b = frame.split()  # 分离成RGB三个通道。。提取R G B分量
                frame = Image.merge('RGB', (b, g, r))  # 合并通道
                # 侦测人脸
                boxes, probs, points = mtcnn.detect(frame, landmarks=True)

                if probs is None or boxes is None:
                    continue
                if len(boxes) > 1 or probs < 0.998:
                    continue
                count += 1

                # Draw boxes and save faces
                # img_draw = frame.copy()
                # draw = ImageDraw.Draw(img_draw)
                # for i, (box, point) in enumerate(zip(boxes, points)):
                #     draw.rectangle(box.tolist(), width=5)
                #     for p in point:
                #         draw.rectangle((p - 10).tolist() + (p + 10).tolist(), width=10)
                #     extract_face(frame, box, save_path='detected_face_{}.png'.format(i))
                # img_draw.save('annotated_faces.png')

                # box是(left, upper, right, lower)
                # 把box变正方形
                for box in boxes:
                    # (l, t, r, b) = box
                    # w = r - l
                    # h = b - t
                    #
                    # size = max(r - l, b - t) * 1.1
                    # xpadding = (size - w) / 2
                    # ypadding = (size - h) / 2
                    #
                    # L = l - xpadding
                    # R = L + size
                    # if L < 0:
                    #     R -= L  # 右边补齐一下，一般左边超了右边不会超
                    #     L = 0
                    # if R > frame.size[0]:
                    #     offset = R - frame.size[0]
                    #     R = frame.size[0]
                    #     L -= offset
                    # T = int(t) - ypadding
                    # B = int(T + size)
                    # if T < 0:
                    #     B -= T
                    #     T = 0
                    # if B > frame.size[1]:
                    #     offset = B - frame.size[1]
                    #     R = frame.size[1]
                    #     L -= offset
                    # box = (L, T, R, B)

                    cropimg = frame.crop(box)
                    yield cropimg
            else:
                print('read error')
            progress(int(count * 100.0 / target))
            if count >= target:
                break
        # 对于不好处理的提前结束
        if trycount >= 50 * target:
            break
    print('取到的张数： ', count, ' trycount:', trycount)
    vc.release()


def image_generator(path_, sec=1):
    """sec 秒取1张"""
    vc = cv2.VideoCapture(path_)
    fps = vc.get(cv2.CAP_PROP_FPS)
    fcount = vc.get(cv2.CAP_PROP_FRAME_COUNT)
    fpsgap = round(fps * sec)
    if fpsgap < 1:
        fpsgap = 1
    print('fps:', fps, ', fpsgap:', fpsgap, ', total frame:', fcount)
    c = 1
    totalframe = 0
    vc.set(cv2.CAP_PROP_POS_FRAMES, totalframe)
    curmin = 5
    while totalframe < fcount:  # 循环读取视频帧
        rval, frame = vc.read()
        if rval:
            # plt.imshow(frame)
            # plt.show()

            yield frame
        else:
            print('read error')
            print(totalframe)
        # 打印时间
        cursec = int(totalframe / fps)
        if cursec / 300 > curmin:
            print("cur min:", cursec / 60.0)
            curmin += 5
        totalframe += fpsgap
        vc.set(cv2.CAP_PROP_POS_FRAMES, totalframe)
    vc.release()


def tensor_to_np(tensor):
    """针对这个模型做的运算，MTCNN中有个fixed_image_standardization，这个是其逆运算"""
    img = tensor.mul(128.0).add(127.5).byte()
    img = img.cpu().numpy().transpose((1, 2, 0))
    return img


def thr_getimg(namepath, videos, savefolder):
    """videos的每个video"""
    for video_name in videos:
        print(video_name)
        videopath = os.path.join(namepath, video_name)
        image_generator_count(videopath)
        for img_face in image_generator_count(videopath, 4):
            # 构造文件名
            uid = ''.join(str(uuid.uuid4()).split('-'))
            savepath = os.path.join(savefolder, uid + '.jpg')
            img_face.save(savepath)

            # cv2.imwrite(savepath, cv2.cvtColor(img_face, cv2.COLOR_RGB2BGR))
            # 一种保存方法
            # img = Image.fromarray(img_face)
            # r, g, b = img.split()  # 分离成RGB三个通道。。提取R G B分量
            # img = Image.merge('RGB', (b, g, r))  # 合并通道
            # img.save(savepath)
            # 另一种保存方法
            # cv2.imencode('.jpg', img_face)[1].tofile(savepath)


def get_video_face(name_path, save_path):
    """
    从video_path的video中获取人脸，将其保存在video_path下的capture_image文件夹下
    每个视频取的张数可指定在函数thr_getimg中
    """
    names = os.listdir(name_path)
    trainroot = os.path.join(save_path, 'train')
    valroot = os.path.join(save_path, 'val')
    testroot = os.path.join(save_path, 'test')
    for name in names:
        namepath = os.path.join(name_path, name)
        videos = os.listdir(namepath)
        name_train_folder = os.path.join(trainroot, name)
        name_val_folder = os.path.join(valroot, name)
        name_test_folder = os.path.join(testroot, name)
        os.makedirs(name_train_folder, exist_ok=True)
        os.makedirs(name_val_folder, exist_ok=True)
        os.makedirs(name_test_folder, exist_ok=True)
        # 移除非video
        tmp = videos
        for video_name in videos:
            if len(video_name.split('.')) == 1:
                tmp.remove(video_name)
                continue
            houzhui = video_name.split('.').pop()
            if houzhui not in ['mp4', 'mkv', 'rmvb', 'avi', 'mov', 'm4s']:
                tmp.remove(video_name)
                continue
        videos = tmp
        # 分为训练和val
        tp = int(len(videos) * 0.7)
        val = int(len(videos) * 0.2)
        train_videos = videos[:tp]
        val_videos = videos[tp:tp + val]
        test_videos = videos[tp + val:]

        thr_getimg(namepath, train_videos, name_train_folder)
        thr_getimg(namepath, val_videos, name_val_folder)
        thr_getimg(namepath, test_videos, name_test_folder)
        # # 创建线程
        # t = threading.Thread(target=thr_getimg, args=(video_name, video_path, savefolder))
        # threads.append(t)

    # for t in threads:
    #     t.start()
    #     print('\nstart thread')
    #     while True:
    #         # 判断正在运行的线程数量,如果小于5则退出while循环,
    #         # 进入for循环启动新的进程.否则就一直在while循环进入死循环
    #         if len(threading.enumerate()) < 3:
    #             time.sleep(1)
    #             break


def get_video_summary(video_path):
    f_save_path = video_path  # 保存图片的上级目录
    videos = os.listdir(video_path)

    for video_name in videos:
        if len(video_name.split('.')) == 1:
            continue
        houzhui = video_name.split('.').pop()
        if houzhui not in ['mp4', 'mkv', 'rmvb', 'avi', 'mov']:
            continue
        file_name = video_name.split('.')[0]
        print(file_name)
        folder_name = os.path.join(f_save_path, file_name)
        os.makedirs(folder_name, exist_ok=True)
        videopath = os.path.join(video_path, video_name)
        imagecount = 0
        # n秒一张截图
        for img in image_generator(videopath, 90):
            savepath = os.path.join(folder_name, file_name + '_' + str(imagecount) + '.jpg')
            cv2.imencode('.jpg', img)[1].tofile(savepath)
            imagecount += 1


if __name__ == '__main__':
    get_video_face(r'H:\48livevideo', r'H:\48faceset')

    # path = r'H:\mlface\wy'
    # get_video_face(path)
    #
    # path = r'H:\mlface\zsy'
    # get_video_face(path)

    # v = r'H:\mlface\jy'
    # get_video_summary(v)
